import sys
import os
import requests
import hashlib
import argparse
import psutil
from PyQt5.QtWidgets import QApplication, QWidget, QVBoxLayout, QLabel, QProgressBar, QMessageBox
from PyQt5.QtCore import QThread, pyqtSignal, pyqtSlot


class UpdateWorker(QThread):
    update_progress = pyqtSignal(int)
    update_message = pyqtSignal(str)
    update_finished = pyqtSignal()

    def __init__(self, version, check_url, original_program, project_url):
        super().__init__()
        self.current_version = version
        self.check_url = check_url
        self.original_program = original_program
        self.project_url = project_url

    def run(self):
        try:
            self.update_message.emit("检查更新中...")
            response = requests.get(self.check_url)
            data = response.json()

            latest_version = data["release"]["tag"]["name"]
            update_url = f"{self.project_url}/releases/download/{latest_version}/update.json"
            self.update_message.emit(f"获取最新版本的文件信息：{latest_version}")
            self.process_update(update_url)

        except Exception as e:
            self.update_message.emit(f"更新失败: {str(e)}")
            self.update_finished.emit()

    def process_update(self, update_url):
        try:
            response = requests.get(update_url)
            if response.status_code == 200:
                update_data = response.json()
                self.update_files(update_data)
            else:
                self.update_message.emit("无法获取更新信息文件")
                self.update_finished.emit()
        except Exception as e:
            self.update_message.emit(f"处理更新时发生错误: {str(e)}")
            self.update_finished.emit()

    def update_files(self, update_data):
        local_files = self.get_local_files_md5()

        for file_path, md5 in update_data.items():
            local_file_md5 = local_files.get(file_path)
            if not local_file_md5 or local_file_md5 != md5:
                self.update_message.emit(f"下载更新文件：{file_path}")
                self.download_and_save_file(file_path)

        self.remove_extra_files(update_data, local_files)
        self.update_finished.emit()

    def get_local_files_md5(self):
        local_files_md5 = {}
        for root, _, files in os.walk("."):
            for file in files:
                file_path = os.path.join(root, file)
                relative_path = os.path.relpath(file_path, ".")
                local_files_md5[relative_path] = self.calculate_md5(file_path)
        return local_files_md5

    def calculate_md5(self, file_path):
        hash_md5 = hashlib.md5()
        with open(file_path, "rb") as f:
            for chunk in iter(lambda: f.read(4096), b""):
                hash_md5.update(chunk)
        return hash_md5.hexdigest()

    def download_and_save_file(self, file_path):
        download_url = f"{self.project_url}/{file_path}"
        response = requests.get(download_url, stream=True)
        os.makedirs(os.path.dirname(file_path), exist_ok=True)
        with open(file_path, 'wb') as file:
            for chunk in response.iter_content(chunk_size=1024):
                file.write(chunk)

    def remove_extra_files(self, update_data, local_files):
        for file_path in local_files.keys():
            if file_path not in update_data:
                self.update_message.emit(f"删除多余文件：{file_path}")
                os.remove(file_path)

    def close_original_program(self):
        try:
            for proc in psutil.process_iter(['name']):
                if proc.info['name'] == self.original_program:
                    proc.kill()
                    proc.wait()
                    break
        except Exception as e:
            self.update_message.emit(f"关闭原程序失败: {str(e)}")

    def start_original_program(self):
        try:
            os.startfile(self.original_program)
        except Exception as e:
            self.update_message.emit(f"启动原程序失败: {str(e)}")


class UpdateProgram(QWidget):
    def __init__(self, version, check_url, original_program):
        super().__init__()
        self.init_ui()
        self.worker = UpdateWorker(version, check_url, original_program, project_url)

        self.worker.update_progress.connect(self.progress.setValue)
        self.worker.update_message.connect(self.label.setText)
        self.worker.update_finished.connect(self.on_update_finished)

        self.worker.start()

    def init_ui(self):
        self.setWindowTitle('更新程序')
        self.layout = QVBoxLayout()

        self.label = QLabel('初始化...')
        self.layout.addWidget(self.label)

        self.progress = QProgressBar()
        self.layout.addWidget(self.progress)

        self.setLayout(self.layout)
        self.resize(300, 100)
        self.show()

    @pyqtSlot()
    def on_update_finished(self):
        self.worker.close_original_program()
        self.worker.start_original_program()
        QMessageBox.information(self, "更新完成", "更新已成功完成。")
        QApplication.quit()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="更新程序")
    parser.add_argument("--version", type=str, required=True, help="当前版本号")
    parser.add_argument("--check_url", type=str, required=True, help="检查更新的URL")
    parser.add_argument("--original_program", type=str, required=True, help="原程序的名称或路径")
    args = parser.parse_args()

    app = QApplication(sys.argv)
    update_prog = UpdateProgram(args.version, args.check_url, args.original_program)
    sys.exit(app.exec_())
